package cloud.daodao.storm.common.util.security;

import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.*;

public class Sm4Util {

    private static final String NAME = "SM4";
    private static final Integer SIZE = 128;
    private static final String ALGORITHM = "SM4/ECB/PKCS7Padding";
    // private static final String ALGORITHM = "SM4/ECB/NoPadding";

    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    public static String key() throws NoSuchAlgorithmException, NoSuchProviderException {
        KeyGenerator keyGenerator = KeyGenerator.getInstance(NAME, BouncyCastleProvider.PROVIDER_NAME);
        keyGenerator.init(SIZE, new SecureRandom());
        SecretKey secretKey = keyGenerator.generateKey();
        byte[] bytes = secretKey.getEncoded();
        return bytesToHex(bytes).toUpperCase();
    }

    public static String encrypt(String key, String iv, String plains) throws NoSuchPaddingException, NoSuchAlgorithmException, NoSuchProviderException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
        if (null == key || key.isEmpty() || null == plains || plains.isEmpty()) {
            return plains;
        }
        Cipher cipher = Cipher.getInstance(ALGORITHM, BouncyCastleProvider.PROVIDER_NAME);
        SecretKeySpec secretKeySpec = new SecretKeySpec(hexToBytes(key), NAME);
        if (null != iv && !iv.isEmpty()) {
            IvParameterSpec ivSpec = new IvParameterSpec(hexToBytes(iv));
            cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivSpec);
        } else {
            cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec);
        }
        cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec);
        byte[] bytes = cipher.doFinal(plains.getBytes(StandardCharsets.UTF_8));
        return bytesToHex(bytes).toUpperCase();
    }

    public static String decrypt(String key, String iv, String ciphers) throws NoSuchPaddingException, NoSuchAlgorithmException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException, InvalidKeyException, NoSuchProviderException {
        if (null == key || key.isEmpty() || null == ciphers || ciphers.isEmpty()) {
            return ciphers;
        }
        Cipher cipher = Cipher.getInstance(ALGORITHM, BouncyCastleProvider.PROVIDER_NAME);
        SecretKeySpec secretKeySpec = new SecretKeySpec(hexToBytes(key), NAME);
        if (null != iv && !iv.isEmpty()) {
            IvParameterSpec ivSpec = new IvParameterSpec(hexToBytes(iv));
            cipher.init(Cipher.DECRYPT_MODE, secretKeySpec, ivSpec);
        } else {
            cipher.init(Cipher.DECRYPT_MODE, secretKeySpec);
        }
        cipher.init(Cipher.DECRYPT_MODE, secretKeySpec);
        byte[] bytes = cipher.doFinal(hexToBytes(ciphers.toLowerCase()));
        return new String(bytes, StandardCharsets.UTF_8);
    }

    private static String bytesToHex(byte[] bytes) {
        StringBuilder sb = new StringBuilder();
        for (byte b : bytes) {
            sb.append(String.format("%02x", b));
        }
        return sb.toString();
    }

    private static byte[] hexToBytes(String hex) {
        byte[] result = new byte[hex.length() / 2];
        int j = 0;
        for (int i = 0; i < hex.length(); i += 2) {
            result[j++] = (byte) Integer.parseInt(hex.substring(i, i + 2), 16);
        }
        return result;
    }

}
